#include "gradient_descent_optimizer.h"
#include "../graph.h"
#include "../variable.h"

GradientDescentOptimizer::GradientDescentOptimizer(Node* node, double alpha) 
    : Optimizer({}), alpha_(alpha) {
    // construct the gradient graph
    Graph* graph = Graph::getGraph();
    graph -> constructGradientNodeds(node);
    // add variable nodes and gradient nodes to nodeIn vector
    auto node2grad = *(graph -> getGradOfVarNodeMap());
    for(auto iter = node2grad.begin(); iter != node2grad.end(); iter++) {
        addNodeIn(iter -> first);
        addNodeIn(iter -> second);
        iter -> first -> addNodeOut(this);
        iter -> second -> addNodeOut(this);
    }
}

void GradientDescentOptimizer::compute(Context* context) {
    auto nodes = *getNodeIn();
    for(int i = 0; i < nodes.size(); i += 2) {
        Variable* var = dynamic_cast<Variable*>(nodes[i]);
//        cout << "update before: " << endl << var -> getValue();
//        cout << "grad : " << endl << context -> getFlowedData(nodes[i + 1]);
        Matrix value = var -> getValue() - alpha_ * context -> getFlowedData(nodes[i + 1]);
        var -> setValue(value);
//        cout << "update after: " << endl << var -> getValue();
    }
}
